import cv2
import time
import torch
import copy
import numpy as np
try:
    from models.experimental import attempt_load
    from utils.general import non_max_suppression, scale_coords, check_img_size
    from utils.datasets import letterbox
    from utils.plots import plot_one_box, colors
    from utils.QRSolver import qrsolver
except:
    from .models.experimental import attempt_load
    from .utils.general import non_max_suppression, scale_coords, check_img_size
    from .utils.datasets import letterbox
    from .utils.plots import plot_one_box, colors
    from .utils.QRSolver import qrsolver

class ToyotaDet(object):
    def __init__(self, weights, device='cuda:0', imgsz=1024):
        self.weights = weights
        self.device = torch.device(device)
        self.model = attempt_load(weights, self.device)
        self.model.half()
        # self.model.eval()
        self.stride = int(self.model.stride.max())  # model stride
        #self.imgsz = check_img_size(imgsz, s=self.stride)
        self.imgsz = imgsz
        self.names = self.model.names

    @torch.no_grad()
    def predict(self, source, plot=False):
        # run model for prediction
        #source = cv2.imread(image_path)
        image = copy.deepcopy(source)
        temp = copy.deepcopy(source)
        source = letterbox(source, new_shape=self.imgsz, stride=self.stride)[0]
        source = source[:, :, ::-1].transpose(2, 0, 1)
        source = np.ascontiguousarray(source)
        source = np.expand_dims(source, 0)
        source = torch.from_numpy(source).to(self.device)
        source = source.half()  # float()
        source /= 255.0
        pred = self.model(source, augment=False)[0]
        pred = non_max_suppression(pred, 0.05, 0.05, None, False, 1000)

        items = dict()
        for i, det in enumerate(pred):
            det[:, :4] = scale_coords(source.shape[2:], det[:, :4], image.shape).round()

            for *xyxy, conf, cls in reversed(det.data.cpu().numpy()):
                c = int(cls)  # integer class
                # label = f'{self.names[c]}'
                label = f'{self.names[c]} {conf:.2f}'
                key = f'{self.names[c]}'
                xmin, ymin, xmax, ymax = int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])

                if key == 'bq' and xmin>1000 and xmax<4000:
                    print('bq')
                    print(ymin, ymax, xmin, xmax)
                    # QRImage = temp[ymin:ymax, xmin:xmax, :]
                    # barcode, tags = qrsolver(QRImage)

                if key in items.keys():
                   
                    items[key].append([xmin, ymin, xmax, ymax, conf])
                else:
                    items[key] = list()
                    items[key].append([xmin, ymin, xmax, ymax, conf])

                if plot:
                    plot_one_box(xyxy, image, label=label, color=colors(c, True), line_thickness=3)

        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        return image, items

if __name__ == '__main__':
    weights = '/media/Harddisk/Users/fan/PROJECT/CDPROJECT/TensorRT/Models/models/model_parts.pt'
    img_path = '/media/Harddisk/Users/fan/PROJECT/CDPROJECT/TensorRT/TestData/C50_FR-R/171118/169.254.7.11.jpg'
    source = cv2.imread(img_path)
    source = cv2.cvtColor(source, cv2.COLOR_BGR2RGB)

    start = time.time()
    det = ToyotaDet(weights=weights)
    finish1 = time.time()
    image, QRImage, items, barcode, tags = det.predict(source, plot=True)
    finish2 = time.time()
    print('Whole time: ', finish2-start)
    print('Det Cost time: ', finish2-finish1)
    cv2.imwrite('QRImage.jpg', QRImage)
    cv2.imwrite('169.254.7.11.jpg', image)
    print(barcode)
    print(tags)


